import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import time

# 修改后的堆排序算法
def heapify(arr, n, i, steps):
    largest = i
    l = 2 * i + 1
    r = 2 * i + 2

    if l < n and arr[i] < arr[l]:
        largest = l

    if r < n and arr[largest] < arr[r]:
        largest = r

    if largest != i:
        steps.append((arr.copy(), i, largest))
        arr[i], arr[largest] = arr[largest], arr[i]
        steps.append((arr.copy(), i, largest))  # 记录排序过程
        heapify(arr, n, largest, steps)

def heap_sort(arr):
    n = len(arr)
    steps = []

    # Build a maxheap
    for i in range(n // 2 - 1, -1, -1):
        heapify(arr, n, i, steps)

    # One by one extract elements
    for i in range(n - 1, 0, -1):
        steps.append((arr.copy(), 0, i))
        arr[i], arr[0] = arr[0], arr[i]
        steps.append((arr.copy(), 0, i))  # 记录排序过程
        heapify(arr, i, 0, steps)

    return steps

# 随机生成一个数组
arr = np.random.randint(0, 100, 20)

# 获取排序步骤
steps = heap_sort(arr.copy())

# 创建动画
fig, ax = plt.subplots()
ax.set_xlim(-1, len(arr))
ax.set_ylim(0, max(arr) + 10)

bar_rects = ax.bar(range(len(arr)), arr, align="edge")

def animate(i):
    current_step = steps[i][0]
    swap_idx1 = steps[i][1]
    swap_idx2 = steps[i][2]

    for idx, (rect, val) in enumerate(zip(bar_rects, current_step)):
        if idx == swap_idx1 or idx == swap_idx2:
            rect.set_color('r')  # 将交换的柱子设为红色
        else:
            rect.set_color('b')  # 其余柱子设为蓝色
        rect.set_height(val)
    time.sleep(0.5)  # 可以适当调整延迟以观察排序过程

    return bar_rects

ani = animation.FuncAnimation(fig, animate, frames=len(steps), interval=50, blit=True)

plt.show()
